# -*- coding: utf-8 -*-

"""
将OD算法模型导出的rworker_match.json文件转换为YAML格式

导出两个文件：
rworker_match_alg.yaml
labeled.yaml

rworker_match_alg.yaml记录了OD算法模型匹配的结果。
labeled.yaml在rworker_match_alg.yaml基础之上，对每一对pair，
增加了 source: algorithm 和 status: unchecked 信息，用于后面的标注工具。

设计的思路：
rworker_match_alg.yaml 记录了OD算法模型匹配的结果，用于后续统计OD算法模型的精度
labeled.yaml记录了标注的结果，通过 source的值 algorithm 和 manual 区分出是算法匹配还是人工查找出来
    通过 status 记录对该匹配对的核实结果
    在标注的过程中，保留right,unchecked和uncertain的记录
    最终标注和检查结束后，只会保留right和uncertain的记录
"""

import os
from datetime import datetime
import json
import oyaml as yaml
import numpy as np

from od.utilities import read_pkg_file


def read_rworker_match_file(match_file_path):
    with open(match_file_path, "r") as match_file:
        matched_pairs = json.load(match_file)
    return matched_pairs


def reformat_rworker_matched_pairs(root, rworker_matched_pairs):
    pairs = dict()
    for i, matched_pair in enumerate(rworker_matched_pairs):
        gallery = {
            "image": F"information/person_tracker_image/{matched_pair['on_lid']:06}",
            "reid_feature": F"information/reid_pkgs/pkg_{matched_pair['on_lid']}.txt",
            "time": datetime.fromtimestamp(
                int(matched_pair['on_time'])).strftime('%Y-%m-%d %H:%M:%S'),
            "location": "",
        }
        query = {
            "image": F"information/person_tracker_image/{matched_pair['off_lid']:06}",
            "reid_feature": F"information/reid_pkgs/pkg_{matched_pair['off_lid']}.txt",
            "time": datetime.fromtimestamp(
                int(matched_pair['off_time'])).strftime('%Y-%m-%d %H:%M:%S'),
            "location": ""
        }
        gallery_feature = read_pkg_file(
            os.path.join(root, gallery["reid_feature"]))["uodr_pkg"]
        query_feature = read_pkg_file(
            os.path.join(root, query["reid_feature"]))["uodr_pkg"]
        distance = 1 - np.dot(np.array(gallery_feature), np.array(query_feature))
        distance = F"{distance.item():.7f}"
        pair = {"gallery": gallery,
                "query": query,
                "distance": distance}
        pairs[F"pair_{i+1:05}"] = pair
    return pairs


def main():
    bus_date_folder = r"E:\data\59092_00E0B453CE0B\2019_3_18"
    rworker_match_file_path = os.path.join(bus_date_folder, "information", "rworker_match.json")

    label_folder = os.path.join(bus_date_folder, "information", "label")
    os.makedirs(label_folder, exist_ok=True)

    rworker_matched_pairs = read_rworker_match_file(rworker_match_file_path)
    pairs = reformat_rworker_matched_pairs(bus_date_folder, rworker_matched_pairs)

    # 保存 rworker_match_alg.yaml
    yaml_file_path = os.path.join(label_folder, "rworker_match_alg.yaml")
    with open(yaml_file_path, "w", encoding="utf-8") as yaml_file:
        yaml.dump(
            {"root": bus_date_folder, "pairs": pairs},
            yaml_file,
            allow_unicode=True,
            default_flow_style=False)
    # 保存 label.yaml
    for pair_name, pair in pairs.items():
        pair["source"] = "algorithm"
        pair["status"] = "unchecked"
    yaml_file_path = os.path.join(label_folder, "od_label.yaml")
    with open(yaml_file_path, "w", encoding="utf-8") as yaml_file:
        yaml.dump(
            {"root": bus_date_folder, "pairs": pairs},
            yaml_file,
            allow_unicode=True,
            default_flow_style=False)


if __name__ == "__main__":
    main()
